{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Acoustic non-linear wave-equation operator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Acoustic isotropic wave equation and stepping rules"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook we show how to derive and implement acoustic isotropic wave-equation operators with simple stepping rules. We start by considering the acoustic isotropic wave equation that can be written as follows:\n",
    "\n",
    "\\begin{equation}\n",
    "\\left[\\frac{1}{\\mathbf{v}^2(\\mathbf{x})}\\frac{\\partial^2}{\\partial t^2} - \\nabla^2 \\right]\\mathbf{p}(\\mathbf{x},t) = \\mathbf{f}(\\mathbf{x},t),\n",
    "\\end{equation}\n",
    "\n",
    "where $\\mathbf{p}$ represents the pressure field, $\\mathbf{f}$ is the forcing term, and $\\mathbf{v}$ is the acoustic velocity. Let's now discretize in time this equation and use a second-order scheme and drop the dependency over the spatial coordinate vector $\\mathbf{x}$. The previous equation now becomes:\n",
    "\n",
    "\\begin{equation}\n",
    "\\frac{\\mathbf{p}(t_{i+1}) -2\\mathbf{p}(t_{i}) +\\mathbf{p}(t_{i-1})}{\\Delta t^2 \\mathbf{v}^2} - \\nabla^2\\mathbf{p}(t_i) = \\mathbf{f}(t_i),\n",
    "\\end{equation}\n",
    "\n",
    "in which the variable $i$ represents index of time sample in the discrete axis with $i \\in  [0,\\dots,N_t]$. Assuming that $\\mathbf{p}(t) = 0$ for $t \\leq 0$, we can write the second-order time derivative as a lower triangular matrix: \n",
    "\n",
    "\\begin{equation}\n",
    "\\mathbf{D}^2_t = \\frac{1}{\\Delta t^{2}}\\begin{bmatrix} 1 & 0 & 0 & \\cdots & 0 & 0 & 0 \\\\\n",
    "                                                 -2 & 1 & 0 & \\cdots & 0 & 0 & 0 \\\\\n",
    "                                                 1 & -2 & 1 & \\cdots & 0 & 0 & 0 \\\\\n",
    "                                                 \\vdots & \\vdots & \\vdots & \\ddots & \\vdots & \\vdots & \\vdots \\\\\n",
    "                                                 0 & 0 & 0 & \\cdots & 1 & 0 & 0 \\\\\n",
    "                                                 0 & 0 & 0 & \\cdots & -2 & 1 & 0 \\\\\n",
    "                                                 0 & 0 & 0 & \\cdots & 1 & -2 & 1\n",
    "             \\end{bmatrix},\n",
    "\\end{equation}\n",
    "\n",
    "where $\\Delta t$ represents the time sampling, giving rise to the following time-stepping rule:\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathbf{p}(t_{i}) =\\Delta t^2 \\mathbf{v}^2\\mathbf{f}(t_{i-1}) + [\\Delta t^2 \\mathbf{v}^2\\nabla^2 +2] \\mathbf{p}(t_{i-1}) - \\mathbf{p}(t_{i-2}),\n",
    "\\end{equation}\n",
    "\n",
    "whose syntax can be simplified by introducing $\\mathbf{A}=\\Delta t^2 \\mathbf{v}^2$ and $\\mathbf{B}=\\mathbf{A}\\nabla^2+2$. Therefore, we can rewrite the previous stepping rule as follows:\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathbf{p}(t_{i}) =\\mathbf{A}\\mathbf{f}(t_{i-1}) + \\mathbf{B} \\mathbf{p}(t_{i-1}) - \\mathbf{p}(t_{i-2}).\n",
    "\\end{equation}\n",
    "\n",
    "Alomomin (2013) states the following: \"Although the previous equation describes the forward time-stepping correctly, its\n",
    "adjoint is not a recursive operator.\". Hence, in their derivation of the adjoint stepping rule they define the $\\mathbf{q}(t_{i})=\\mathbf{A}\\mathbf{f}(t_{i})$, which changes the forward stepping rule as follows:\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathbf{p}(t_{i}) =\\mathbf{q}(t_{i-1}) + \\mathbf{B} \\mathbf{p}(t_{i-1}) - \\mathbf{p}(t_{i-2}),\n",
    "\\end{equation}\n",
    "\n",
    "which allows them to write the following adjoint stepping rule:\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathbf{q}(t_{i}) =\\mathbf{p}(t_{i+1}) + \\mathbf{B}^{*} \\mathbf{q}(t_{i+1}) - \\mathbf{q}(t_{i+2}),\n",
    "\\end{equation}\n",
    "\n",
    "where adjoint operator of $\\mathbf{B}$ is given by the following expression:\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathbf{B}^{*} = \\nabla^{2*}\\mathbf{A}^{*} + 2 = \\nabla^{2*}\\mathbf{A} + 2,\n",
    "\\end{equation}\n",
    "\n",
    "which shows that it is necessary to apply the operator $\\mathbf{A}$ before computing the adjoint Laplacian operator $\\nabla^{2*}$. Even though this simple change of operation order seems a small detail, it actually results in a more complex adjoint operator and possibly in a less efficient implementation. In the next section we show how to avoid this issue and implement a simpler stepping rule that avoids the usage of the operator $\\mathbf{B}^{*}$ altogether."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A different adjoint stepping rule"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's go back to the time-discrete version of the wave-equation introduced in the previous section. Using the operator $\\mathbf{D}^2_t$, we can write:\n",
    "\n",
    "\\begin{equation}\n",
    "\\left[\\frac{1}{\\mathbf{v}^2}\\mathbf{D}^2_t - \\nabla^2\\right]\\mathbf{p} = \\mathbf{F} \\mathbf{p} = \\mathbf{f},\n",
    "\\end{equation}\n",
    "\n",
    "where $\\mathbf{p}=[\\mathbf{p}(t_0),\\dots,\\mathbf{p}(t_{N_t})]^{*}$ and $\\mathbf{f}=[\\mathbf{f}(t_0),\\dots,\\mathbf{f}(t_{N_t})]^{*}$ . The previous section showed that this linear system can be solved by Gaussian substitution to find $\\mathbf{p} = \\mathbf{F}^{-1}\\mathbf{f}$. To find the adjoint $\\mathbf{F}^{*}$ and its inverse operator let's write the following expression:\n",
    "\n",
    "\\begin{equation}\n",
    " \\mathbf{F}^{*} \\mathbf{f} = \\left[\\frac{1}{\\mathbf{v}^2}\\mathbf{D}^2_t - \\nabla^2\\right]^{*} \\mathbf{f} = \\left[\\frac{1}{\\mathbf{v}^2}\\mathbf{D}^{2*}_t - \\nabla^{2*}\\right] \\mathbf{f} = \\mathbf{p},\n",
    "\\end{equation}\n",
    "\n",
    "where we used the fact that $\\frac{1}{\\mathbf{v}^2}$ commutes with the derivative operator $\\mathbf{D}^{2*}_t$. This equation shows that the only difference from the forward equation is given by the way we solve it, since in the adjoint equation we have a upper triangular system, and the presence of $\\nabla^{2*}$ as opposed to $\\nabla^{2}$, which can be removed if $\\nabla^{2*}=\\nabla^{2}$\n",
    "Again, if we solve the adjoint wave equation by Gaussian substitution, we can write the adjoint stepping rule as follows:\n",
    "\n",
    "\\begin{equation}\n",
    "\\mathbf{f}(t_{i}) =\\mathbf{A}\\mathbf{p}(t_{i+1}) + \\mathbf{B} \\mathbf{f}(t_{i+1}) - \\mathbf{f}(t_{i+2}),\n",
    "\\end{equation}\n",
    "\n",
    "where we assume $\\nabla^{2}$ to be self-adjoint, assumption often fulfilled unless particular boundary conditions are employed (e.g., free-surface boundary). Beside the change in time indices, this stepping rule is computationally equivalent to the one present in the forward wave-equation operator. Let's test if this derivation passes the dot-product test. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Numerical test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start by importing the necessary modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING! DATAPATH not found. The folder /tmp will be used to write binary files\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "sys.path.append(\"/Users/ettorebiondi/PycharmProjects/python-solver/GenericSolver/python\")\n",
    "import pyVector as Vec\n",
    "import pyOperator as Op\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, lets define the wave-equation operator that makes use of the stepping rules described in the previous section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AcousticWaveEquation(Op.Operator):\n",
    "\tdef __init__(self,wavefield,vel,dt,ds):\n",
    "\t\tself.setDomainRange(wavefield,wavefield)\n",
    "\t\tself.vel = vel\n",
    "\t\tself.dt = dt\n",
    "\t\tself.dt2 = dt*dt\n",
    "\t\tself.ds2_inv = 1.0/(ds*ds)\n",
    "\t\tself.vel2dt = self.vel.getNdArray()*self.vel.getNdArray()*self.dt2\n",
    "\n",
    "\tdef forward(self,add,model,data):\n",
    "\t\tself.checkDomainRange(model,data)\n",
    "\t\tif not add:\n",
    "\t\t\tdata.zero()\n",
    "\t\tm_arr = model.clone().getNdArray()\n",
    "\t\t# Scaling input wavefield\n",
    "\t\tm_arr *= self.vel2dt\n",
    "        # Getting number of samples per axis\n",
    "\t\tdata_tmp = data.clone().zero()\n",
    "\t\td_arr = data_tmp.getNdArray()\n",
    "\t\tnt = m_arr.shape[0]\n",
    "\t\tnx = m_arr.shape[1]\n",
    "\t\tnz = m_arr.shape[2]\n",
    "        # Imposing spatial boundary conditions\n",
    "\t\tm_arr[:,0,:]=0.0\n",
    "\t\tm_arr[:,-1,:]=0.0\n",
    "\t\tm_arr[:,:,0]=0.0\n",
    "\t\tm_arr[:,:,-1]=0.0\n",
    "        # Stepping \n",
    "\t\tfor it in range(nt):\n",
    "\t\t\tif it > 1:\n",
    "\t\t\t\tfor idx in range(1,nx-1):\n",
    "\t\t\t\t\tfor idz in range(1,nz-1):\n",
    "\t\t\t\t\t\td_arr[it,idx,idz] += (d_arr[it-1,idx,idz-1]+d_arr[it-1,idx,idz+1]-4.0*d_arr[it-1,idx,idz]+d_arr[it-1,idx-1,idz]+d_arr[it-1,idx+1,idz])*self.ds2_inv*self.vel2dt[idx,idz] \\\n",
    "                                              + 2.0*d_arr[it-1,idx,idz] - d_arr[it-2,idx,idz] + m_arr[it-1,idx,idz]\n",
    "\t\t\telif it == 1:\n",
    "\t\t\t\tfor idx in range(1,nx-1):\n",
    "\t\t\t\t\tfor idz in range(1,nz-1):\n",
    "\t\t\t\t\t\td_arr[it,idx,idz] += (d_arr[it-1,idx,idz-1]+d_arr[it-1,idx,idz+1]-4.0*d_arr[it-1,idx,idz]+d_arr[it-1,idx-1,idz]+d_arr[it-1,idx+1,idz])*self.ds2_inv*self.vel2dt[idx,idz] \\\n",
    "                                              + 2.0*d_arr[it-1,idx,idz] + m_arr[it-1,idx,idz]\n",
    "\t\tdata.scaleAdd(data_tmp)\n",
    "\n",
    "\n",
    "\tdef adjoint(self,add,model,data):\n",
    "\t\tself.checkDomainRange(model,data)\n",
    "\t\tif not add:\n",
    "\t\t\tmodel.zero()\n",
    "\t\td_arr = data.clone().getNdArray()\n",
    "\t\tmodel_tmp = model.clone().zero()\n",
    "\t\tm_arr = model_tmp.getNdArray()\n",
    "        # Scaling input wavefield\n",
    "\t\td_arr *= self.vel2dt\n",
    "        # Getting number of samples per axis\n",
    "\t\tnt = m_arr.shape[0]\n",
    "\t\tnx = m_arr.shape[1]\n",
    "\t\tnz = m_arr.shape[2]\n",
    "        # Imposing spatial boundary conditions\n",
    "\t\td_arr[:,0,:]=0.0\n",
    "\t\td_arr[:,-1,:]=0.0\n",
    "\t\td_arr[:,:,0]=0.0\n",
    "\t\td_arr[:,:,-1]=0.0\n",
    "        # Stepping \n",
    "\t\tfor it in range(nt,-1,-1):\n",
    "\t\t\tif it < nt-2:\n",
    "\t\t\t\tfor idx in range(1,nx-1):\n",
    "\t\t\t\t\tfor idz in range(1,nz-1):\n",
    "\t\t\t\t\t\tm_arr[it,idx,idz] += (m_arr[it+1,idx,idz-1]+m_arr[it+1,idx,idz+1]-4.0*m_arr[it+1,idx,idz]+m_arr[it+1,idx-1,idz]+m_arr[it+1,idx+1,idz])*self.ds2_inv*self.vel2dt[idx,idz] \\\n",
    "                                              + 2.0*m_arr[it+1,idx,idz] - m_arr[it+2,idx,idz] + d_arr[it+1,idx,idz]\n",
    "\t\t\telif it == nt-2:\n",
    "\t\t\t\tfor idx in range(1,nx-1):\n",
    "\t\t\t\t\tfor idz in range(1,nz-1):\n",
    "\t\t\t\t\t\tm_arr[it,idx,idz] += (m_arr[it+1,idx,idz-1]+m_arr[it+1,idx,idz+1]-4.0*m_arr[it+1,idx,idz]+m_arr[it+1,idx-1,idz]+m_arr[it+1,idx+1,idz])*self.ds2_inv*self.vel2dt[idx,idz] \\\n",
    "                                              + 2.0*m_arr[it+1,idx,idz] + d_arr[it+1,idx,idz]\n",
    "\t\tmodel.scaleAdd(model_tmp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's instantiate an acoustic wave-equation operator and test its adjointness."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing wave-equation operator\n",
      "Dot-product test of forward and adjoint operators\n",
      "-------------------------------------------------\n",
      "Applying forward operator add=False\n",
      " Runs in: 7.381800889968872 seconds\n",
      "Applying adjoint operator add=False\n",
      " Runs in: 7.315475940704346 seconds\n",
      "Dot products add=False: domain=2.089362e-02 range=2.089362e-02 \n",
      "Absolute error: 6.245005e-17\n",
      "Relative error: 2.988953e-15 \n",
      "\n",
      "Applying forward operator add=True\n",
      " Runs in: 7.108118057250977 seconds\n",
      "Applying adjoint operator add=True\n",
      " Runs in: 7.236320972442627 seconds\n",
      "Dot products add=True: domain=4.178723e-02 range=4.178723e-02 \n",
      "Absolute error: 1.249001e-16\n",
      "Relative error: 2.988953e-15 \n",
      "\n",
      "-------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "nx=150\n",
    "nz=210\n",
    "nt=60\n",
    "wav = Vec.vectorIC(np.zeros((nt,nx,nz)))\n",
    "vel = Vec.vectorIC(np.zeros((nx,nz)))\n",
    "vel.set(2.0)\n",
    "vel.getNdArray()[:,20:] = 2.5\n",
    "vel.getNdArray()[:,:] += np.random.rand(nx,nz)*0.1\n",
    "ds = 0.02\n",
    "dt = 0.001\n",
    "print(\"Testing wave-equation operator\")\n",
    "prop = AcousticWaveEquation(wav,vel,dt,ds)\n",
    "prop.dotTest(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ali Almomin, 2013, Accurate implementation of two-way wave-equation operators, SEP report 149"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
